import sys
sys.path.append(".")

import numpy as np
import numpy.random as npr
import matplotlib as mpl
import matplotlib.pyplot as plt
import copy, string, itertools, bisect, math, time


#! Parameters:
F  = 15   # fecundity = number of offspring per female
Na = F+1  # number of alleles per locus
transmission_rate = 0.  # transmission rate of the male-killer
lethality = 1.  # killing efficiency
reallocation_efficiency = 1.
mutation_rate = 1e-5

boolval = {'uninfected': 0, 'infected': 1}

alleles = ["a{}".format(i) for i in range(Na)]

phenotypic_value = dict()
for i in range(Na):
    a = alleles[i]
    phenotypic_value[a] = i


def randchoice(seq):
    return seq[npr.randint(len(seq))]
    
def random_genotype():
    return [randchoice(alleles), randchoice(alleles)]

    
class WeightedRandomGenerator(object):
    def __init__(self, weights):
        self.totals = []
        running_total = 0

        for w in weights:
            running_total += w
            self.totals.append(running_total)

    def next(self):
        rnd = npr.random() * self.totals[-1]
        return bisect.bisect_right(self.totals, rnd)

    def __call__(self):
        return self.next()

class WeightedRandomGeneratorWithoutReplacement(object):
    def __init__(self, weights):
        self.totals = []
        self.selected = []
        running_total = 0

        for w in weights:
            running_total += w
            self.totals.append(running_total)

    def next(self):
        if len(self.totals) == len(self.selected):
            raise ValueError, "generator is empty"
        rnd = npr.random() * self.totals[-1]
        i = bisect.bisect_right(self.totals, rnd)
        if i not in self.selected:
            self.selected.append(i)
            return i
        else:
            return self.next()
        
    def __call__(self):
        return self.next()

def shift_mutate(allele):
    i = alleles.index(allele)
    if npr.random() < 0.5:                # shift left
        if not i == 0:
            i -= 1
    else:                               # shift right
        if not i == F:
            i += 1
    return alleles[i]

def jump_mutate(allele):
    mutant_alleles = alleles[:]
    mutant_alleles.remove(allele)
    return randchoice( mutant_alleles )
        
def mutate(allele, mode='shift'):
    if mode=='jump':
        return jump_mutate(allele)
    else:
        return shift_mutate(allele)

def fitness_compensation(x, F, b, rmin=None):
    """x    - number of killed sons
       F    - total clutch size
       rmin - resource threshold to viability
       b    - reallocation efficiency"""
    if not rmin:
        rmin = 1./(10*F)            # default value: optimal resource level is 10 times the threshold
    a = (1-F*rmin)*(F-x)
    numerator   = a + b*x
    denominator = a + b*x*F*rmin
    return numerator / denominator

class Individual(object):
    def __init__(self, genotype, infection_status, sex=None, fitness=1.):
        self.genotype = genotype
        self.infection_status = infection_status
        self.infected = boolval[infection_status]
        self.sex = sex
        self.fitness = fitness
        self.phenotype = self.get_phenotypic_value()
        
    def __str__(self):
        return self.show()
    
    def show(self):
        return """{}: {}
    phenotype: {}
    fitness: {}
    genotype: {}
    """.format( self.sex, self.infection_status, self.phenotype, self.fitness, self.genotype2string() )
    
    def genotype2string(self):
        return '{}  ||  {}'.format(self.genotype[0],self.genotype[1])
        
    def get_phenotypic_value(self):
        """Take the averag."""
        a1,a2 = self.genotype
        return ( phenotypic_value[a1] + phenotypic_value[a2] ) / 2.
    
    def update_fitness(self, factor):
        self.fitness *= factor
    
    def gamete(self):
        mutation = 0
        g = randchoice( self.genotype )
        # mutation:
        if npr.random() < mutation_rate:
            g = mutate(g)
            mutation = 1
        return g, mutation
        
        

class Female(Individual):
    def __init__(self, genotype, infection_status, fitness=1.):
        Individual.__init__(self, genotype, infection_status, sex='female', fitness=fitness)
        self.sex_ratio = math.floor(self.phenotype)
        self.fecundity = F
    
    def show(self):
        return """{}: {}
    sex ratio: {}
    fecundity: {}
    fitness: {}
    genotype: {}
    """.format( self.sex, self.infection_status, self.sex_ratio, self.fecundity, self.fitness, self.genotype2string() )
    
    def produce_brood_with(self, male):
        mut_count = 0
        daughters = []
        sons = []
        for i in range(1,self.fecundity+1):
            # offspring genotype:
            egg, m1   = self.gamete()
            sperm, m2 = male.gamete()
            mut_count += m1+m2
            # offspring infection status:
            offspring_status = 'uninfected'
            if self.infected:
                if npr.random() < transmission_rate:               # random transmission of infection
                    offspring_status = 'infected'
            # offspring sex:
            if i <= self.sex_ratio:
                sons.append( Male([egg,sperm], offspring_status) )
            else:
                daughters.append( Female([egg,sperm], offspring_status) )
        if self.infected:
            sons,daughters = self.male_killing(sons,daughters)
        male.number_of_matings += 1
        return sons, daughters, mut_count
    
    def male_killing(self, sons, daughters):
        n = len(sons)
        sons[:] = (s for s in sons if s.infected and not npr.random()<lethality)
        body_count = n - len(sons)
        for offspring in sons+daughters:
            offspring.update_fitness( fitness_compensation(body_count, self.fecundity, reallocation_efficiency, rmin=None) )
        return sons,daughters
    
class Male(Individual):
    def __init__(self, genotype, infection_status, fitness=1.):
        Individual.__init__(self, genotype, infection_status, sex='male', fitness=fitness)
        self.number_of_matings = 0
        
    def show(self):
        return """{}: {}
    number of matings: {}
    fitness: {}
    genotype: {}
    """.format( self.sex, self.infection_status, self.number_of_matings, self.fitness, self.genotype2string() )
    
def random_individual():
    if npr.random() < 0.5:
        status = 'uninfected'
    else:
        status = 'infected'
    if npr.random() < 0.5:
        return Male(random_genotype(), status)
    else:
        return Female(random_genotype(), status)
        
def random_population(N):
    females = []
    males = []
    for i in range(N):
        individual = random_individual()
        if individual.sex == 'male':
            males.append(individual)
        else:
            females.append(individual)
    return Population(females, males, N)
    
def uniform_population():
    gnum = 1
    N = 2*gnum**Na*Na
    females = []
    males = []
    for i in range(Na):
        for j in range(Na):
            for k in range(gnum):
                if npr.random() < 0.5:
                    status = 'uninfected'
                else:
                    status = 'infected'
                male = Male([alleles[i],alleles[j]],status)
                if npr.random() < 0.5:
                    status = 'uninfected'
                else:
                    status = 'infected'
                female = Female([alleles[i],alleles[j]],status)
                males.append(male)
                females.append(female)
    return Population(females,males)
            
    
    
class Population(object):
                              
    def __init__(self, females, males, N=None, generation=0):
        self.females = females
        self.males = males
        if not N:
            self.size = len(females+males)
        else:
            self.size = N
        self.generation = generation
        record_keys = ['generation','n_females','n_males','mk','mk_females','mk_males','n_mutations','average_sex_ratio','allele_count','genotype_count']
        self.record = {}
        for k in record_keys:
            self.record[k] = []
        self.update_record(mut_count=0)
    
    
    def allele_count(self, g=None):
        if g == None:
            g = self.genotype_count()
        return np.sum(g,0) + np.sum(g,1)
        #        acount = [0 for a in alleles]           # allele count
        #        for i in self.females+self.males:
        #            for a in i.genotype:
        #                acount[ alleles.index(a) ] += 1
        #        return acount
    
    def genotype_count(self):
        """'2D allele count"""
        gcount = np.zeros((Na,Na), dtype=int)
        for i in self.females+self.males:
            a1,a2 = i.genotype
            gcount[alleles.index(a1),alleles.index(a2)] += 1
        return gcount
    
    def asr(self):
        """Calculate average sex ratio (asr) among females."""
        value = 0.
        for f in self.females:
            value += f.sex_ratio
        return value/len(self.females)
    
    def mk(self):
        nF = nM = 0
        for f in self.females:
            if f.infected:
                nF += 1
        for m in self.males:
            if m.infected:
                nM += 1
        return nM, nF
    
    def scan(self):
        sr = 0.
        nF = nM = 0
        acount = [0 for a in alleles]           # allele count
        gcount = np.zeros((Na,Na), dtype=int)   # genotype count
        for f in self.females:
            sr += f.sex_ratio
            if f.infected:
                nF += 1
            a1,a2 = f.genotype
            gcount[ alleles.index(a1), alleles.index(a2) ] += 1
        sr /= len(self.females)
        for m in self.males:
            if m.infected:
                nM += 1
            a1,a2 = m.genotype
            gcount[ alleles.index(a1), alleles.index(a2) ] += 1
        acount = self.allele_count(gcount)
        return sr, nF, nM, acount, gcount            
    
    def update_record(self, mut_count):
        self.record['generation'].append( self.generation )
        self.record['n_males'].append( len(self.males) )
        self.record['n_females'].append( len(self.females) )
        self.record['n_mutations'].append( mut_count )
        sr, mkF, mkM, acount, gcount = self.scan()
        self.record['average_sex_ratio'].append( sr )
        self.record['mk_females'].append( mkF )
        self.record['mk_males'].append( mkM )
        self.record['mk'].append( mkF+mkM )
        self.record['allele_count'].append( acount )
        self.record['genotype_count'].append( gcount )
        
    def is_polymorphic(self):
        acount = self.record['allele_count'][-1]         # most current allele count
        maxcount = 2*self.size
        return maxcount not in acount
        
    def produce_offspring(self):
        mut_count = 0
        if self.males == []:
            raise ValueError, 'population extinct due to lack of males'
        males = self.males[:]
        npr.shuffle(males)
        females = self.females[:]
        npr.shuffle(females)
        temp_pool = []
        offspring = []

        for f in females:
            try:
                m = males.pop()
            except IndexError:
                males = temp_pool[:]
                npr.shuffle(males)
                temp_pool = []
                m = males.pop()
            sons,daughters, mut = f.produce_brood_with(m)
            mut_count += mut
            temp_pool.append(m)
            offspring.extend(sons+daughters)
        return offspring, mut_count
        
    def choose_new_generation(self, offspring):
        females = []
        males = []
        weights = []
        for o in offspring:
            weights.append(o.fitness)
        wrg = WeightedRandomGeneratorWithoutReplacement(weights)
        for i in range(self.size):
            individual = offspring[ wrg() ]
            if individual.sex == 'female':
                females.append(individual)
            else:
                males.append(individual)
        return males, females
    
    def next(self):
        offspring, mut_count = self.produce_offspring()
        self.males, self.females = self.choose_new_generation(offspring)
        self.generation += 1
        self.update_record(mut_count)
    
    def run(self, n, force=False, timing=False):
        """If force is False then n is maximal generation count but iteration stops if the population
        is no longer polymorphic."""
        if timing:
            g0 = self.generation
            t0 = time.time()
        if force:
            for i in range(n):
                self.next()
        else:
            g = 0.
            while self.is_polymorphic() and g<n:
                self.next()
                g += 1
        if timing:
            print "Iterated {} generations in {:.2f} seconds".format(self.generation-g0, time.time()-t0)
            
    def visualize(self, plots=['asr','nM','mk','mut','acount','inds']):
        labels = dict(asr=dict(rec='average_sex_ratio',xlabel='generation',ylabel='average sex ratio',style='k-',ylims=[0,F]), \
                      nM=dict(rec='n_males',xlabel='generation',ylabel='number of males',style='g-',ylims=[0,self.size]), \
                      mk=dict(rec='mk',xlabel='generation',ylabel='MK infected',style='r-',ylims=[0,self.size]), \
                      mut=dict(rec='n_mutations',xlabel='generation',ylabel='number of mutations',style='m-'), \
                      acount=dict(rec='allele_count',xlabel='generation',ylabel='alleles'), \
                      inds=dict(rec='individuals',xlabel='allele 1',ylabel='allele 2') )
        if not isinstance(plots, list):
            plots = list(plots)
        if 'inds' in plots:
            plots.remove('inds')
        figs = []
        g = np.array(self.record['generation'])
        for p in plots:
            f = plt.figure()
            ax = f.add_subplot(111)
            figs.append(f)
            rec = labels[p]['rec']
            xl = labels[p]['xlabel']
            yl = labels[p]['ylabel']
            
            a = np.array(self.record[rec])
            if p == 'mut':
                a = np.cumsum(a)
                labels[p]['ylims']= [0,np.max(a)]
            if p == 'acount':
                a = np.transpose( np.array(a,dtype=float) )
                a /= (2*self.size)                         # normalize (so that values range between 0 and 1)
                ma = np.ma.array(a,mask=(a==0))            # mask zero values
                cmap = mpl.cm.jet                          # use jet color scheme
                cmap.set_bad('w')                          # use white for zero values
                im = ax.imshow(ma,cmap=cmap,vmin=0.,vmax=1.,origin='lower',interpolation='nearest',aspect='auto')
                cb = plt.colorbar(im)
                cb.set_label('allele frequency')
            elif p == 'inds':
                pass
            else:
                ylims = labels[p]['ylims']
                s = labels[p]['style']
                ax.plot(g,a,s)
                ax.set_ylim(*ylims)
            ax.set_xlabel(xl)
            ax.set_ylabel(yl)
        return plots,figs
        
    def genotype_show(self, g=-1):
        f = plt.figure()
        ax = f.add_subplot(111)
        a = np.array( self.record['genotype_count'][g], dtype=float )
        a /= self.size                                      # normalize
        ma = np.ma.array(a, mask=(a==0))
        cmap = mpl.cm.jet
        cmap.set_bad('w')
        im = ax.imshow(ma,cmap=cmap,vmin=0.,vmax=1.,origin='lower',interpolation='nearest',aspect='auto')
        cb = plt.colorbar(im)
        cb.set_label('frequency')
        ax.set_xlabel('chromosome 1')
        ax.set_ylabel('chromosome 2')
        
    
if __name__ == '__main__':
    #    npr.RandomState(21112011)
    npr.seed(21112011)
    
    #    pop = random_population(200)
    pop = uniform_population()
    #    females = [Female(['a1','a2'],'uninfected') for i in range(30)]
    #    males = [Male(['a1','a2'],'uninfected') for i in range(30)]
    #    pop = Population(females,males)
    pop.run(100, force=True, timing=True)
    
    plots,figs = pop.visualize()    
    plt.show()
